use std::rc::Rc;

use common::IndexType;
use itertools::Either;
use rand::distr::Uniform;
use readyset_client::internal::Index;
use readyset_client::KeyComparison;
use readyset_data::{Bound, DfValue, RangeBounds};
use readyset_util::SizeOf;
use vec1::Vec1;

use crate::keyed_state::KeyedState;
use crate::mk_key::MakeKey;
use crate::{LookupResult, PointKey, RangeKey, RangeLookupResult, RecordResult, Row, Rows};

/// A single index of a [`MemoryState`].
///
/// SingleState, which is a wrapper around [`KeyedState`] with extra information about keys and
/// partial, are associative maps from some subset of the columns in a row, to those rows
/// themselves. Internally, as specified by the [`Index`] passed at construction, KeyedState (and
/// hence SingleState) can be backed by either a [`BTreeMap`](std::collections::BTreeMap) or an
/// [`IndexMap`](indexmap::IndexMap) (which is similar to a [`HashMap`](std::collections::HashMap)).
///
/// Any operations on a SingleState that are unsupported by the index type, such as inserting or
/// looking up ranges in a [`HashMap`](readyset_client::IndexType::HashMap) index, will panic, as
/// mixing up index types is an unrecoverable violation of a broad system invariant.
pub(super) struct SingleState {
    /// The column indices and index type that this index is keyed on
    index: Index,

    /// The map containing the state itself.
    ///
    /// Invariant: The length of keys supported by `state` must be equal to `key.len()`
    state: KeyedState,

    /// Is this state partial?
    ///
    /// If this is `false`, lookups (via [`lookup`] or [`lookup_range`]) can never return misses
    partial: bool,

    /// Denormalized cache of the number of rows stored in this index
    row_count: usize,
}

impl SingleState {
    /// Construct a new, empty [`SingleState`] for the given `index`. If `partial`
    pub(super) fn new(index: Index, partial: bool) -> Self {
        let mut state = KeyedState::from(&index);
        if !partial && index.index_type == IndexType::BTreeMap {
            // For fully materialized indices, we never miss - so mark that the full range of keys
            // has been filled.
            state.insert_full_range();
        }
        Self {
            index,
            state,
            partial,
            row_count: 0,
        }
    }

    /// Inserts the given row, or returns false if a hole was encountered (and the record hence
    /// not inserted).
    pub(super) fn insert_row(&mut self, row: Row) -> bool {
        let added = self.state.insert(&self.index.columns, row, self.partial);
        if added {
            self.row_count += 1;
        }
        added
    }

    /// Attempt to remove row `r`.
    pub(super) fn remove_row(&mut self, r: &[DfValue], hit: &mut bool) -> Option<Row> {
        let row = self.state.remove(&self.index.columns, r, Some(hit));
        if row.is_some() {
            self.row_count = self.row_count.saturating_sub(1);
        }
        row
    }

    fn mark_point_filled(&mut self, key: Vec1<DfValue>) {
        let mut key = key.into_iter();
        let replaced = match self.state {
            KeyedState::AllRows(_) => {
                // No-op (can't be partial)
                return;
            }
            KeyedState::SingleBTree(ref mut map) => {
                assert_eq!(key.len(), 1);
                map.insert(key.next().unwrap(), Rows::default())
            }
            KeyedState::DoubleBTree(ref mut map) => {
                assert_eq!(key.len(), 2);
                map.insert((key.next().unwrap(), key.next().unwrap()), Rows::default())
            }
            KeyedState::TriBTree(ref mut map) => {
                assert_eq!(key.len(), 3);
                map.insert(
                    (
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                    ),
                    Rows::default(),
                )
            }
            KeyedState::QuadBTree(ref mut map) => {
                assert_eq!(key.len(), 4);
                map.insert(
                    (
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                    ),
                    Rows::default(),
                )
            }
            KeyedState::QuinBTree(ref mut map) => {
                assert_eq!(key.len(), 5);
                map.insert(
                    (
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                    ),
                    Rows::default(),
                )
            }
            KeyedState::SexBTree(ref mut map) => {
                assert_eq!(key.len(), 6);
                map.insert(
                    (
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                    ),
                    Rows::default(),
                )
            }
            KeyedState::MultiBTree(ref mut map, len) => {
                // I hope LLVM optimizes away the unnecessary into_iter() -> collect()
                assert_eq!(key.len(), len);
                map.insert(key.collect(), Rows::default())
            }
            KeyedState::MultiHash(ref mut map, len) => {
                assert_eq!(key.len(), len);
                map.insert(key.collect(), Rows::default())
            }
            KeyedState::SingleHash(ref mut map) => {
                assert_eq!(key.len(), 1);
                map.insert(key.next().unwrap(), Rows::default())
            }
            KeyedState::DoubleHash(ref mut map) => {
                assert_eq!(key.len(), 2);
                map.insert((key.next().unwrap(), key.next().unwrap()), Rows::default())
            }
            KeyedState::TriHash(ref mut map) => {
                assert_eq!(key.len(), 3);
                map.insert(
                    (
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                    ),
                    Rows::default(),
                )
            }
            KeyedState::QuadHash(ref mut map) => {
                assert_eq!(key.len(), 4);
                map.insert(
                    (
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                    ),
                    Rows::default(),
                )
            }
            KeyedState::QuinHash(ref mut map) => {
                assert_eq!(key.len(), 5);
                map.insert(
                    (
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                    ),
                    Rows::default(),
                )
            }
            KeyedState::SexHash(ref mut map) => {
                assert_eq!(key.len(), 6);
                map.insert(
                    (
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                        key.next().unwrap(),
                    ),
                    Rows::default(),
                )
            }
        };
        assert!(replaced.is_none());
    }

    fn mark_range_filled(&mut self, range: (Bound<Vec1<DfValue>>, Bound<Vec1<DfValue>>)) {
        self.state.insert_range(range);
    }

    /// Marks the given key `filled` and returns the amount of memory bytes used to store
    /// the state.
    pub(super) fn mark_filled(&mut self, key: KeyComparison) {
        match key {
            KeyComparison::Equal(k) => self.mark_point_filled(k),
            KeyComparison::Range(range) => self.mark_range_filled(range),
        }
    }

    /// Remove all rows for the given `key`, and return the amount of memory freed in bytes
    ///
    /// # Panics
    ///
    /// * Panics if the `key` is a range, but the underlying KeyedState is backed by a HashMap
    /// * Panics if the underlying index has no columns
    pub(super) fn mark_hole(&mut self, key: &KeyComparison) -> usize {
        let removed: Box<dyn Iterator<Item = (Row, usize)>> = match key {
            KeyComparison::Equal(key) => match self.state {
                KeyedState::AllRows(_) => panic!("Empty-column index cannot be partial"),
                KeyedState::SingleBTree(ref mut m) => {
                    Box::new(m.remove(&(key[0])).into_iter().flatten())
                }
                KeyedState::DoubleBTree(ref mut m) => {
                    Box::new(m.remove(&MakeKey::from_key(key)).into_iter().flatten())
                }
                KeyedState::TriBTree(ref mut m) => {
                    Box::new(m.remove(&MakeKey::from_key(key)).into_iter().flatten())
                }
                KeyedState::QuadBTree(ref mut m) => {
                    Box::new(m.remove(&MakeKey::from_key(key)).into_iter().flatten())
                }
                KeyedState::QuinBTree(ref mut m) => {
                    Box::new(m.remove(&MakeKey::from_key(key)).into_iter().flatten())
                }
                KeyedState::SexBTree(ref mut m) => {
                    Box::new(m.remove(&MakeKey::from_key(key)).into_iter().flatten())
                }
                KeyedState::MultiBTree(ref mut m, len) => {
                    debug_assert_eq!(key.len(), len);
                    Box::new(m.remove(key.as_vec()).into_iter().flatten())
                }
                KeyedState::MultiHash(ref mut m, len) => {
                    debug_assert_eq!(key.len(), len);
                    Box::new(m.swap_remove(key.as_vec()).into_iter().flatten())
                }
                KeyedState::SingleHash(ref mut m) => {
                    Box::new(m.swap_remove(&(key[0])).into_iter().flatten())
                }
                KeyedState::DoubleHash(ref mut m) => Box::new(
                    m.swap_remove::<(DfValue, _)>(&MakeKey::from_key(key))
                        .into_iter()
                        .flatten(),
                ),
                KeyedState::TriHash(ref mut m) => Box::new(
                    m.swap_remove::<(DfValue, _, _)>(&MakeKey::from_key(key))
                        .into_iter()
                        .flatten(),
                ),
                KeyedState::QuadHash(ref mut m) => Box::new(
                    m.swap_remove::<(DfValue, _, _, _)>(&MakeKey::from_key(key))
                        .into_iter()
                        .flatten(),
                ),
                KeyedState::QuinHash(ref mut m) => Box::new(
                    m.swap_remove::<(DfValue, _, _, _, _)>(&MakeKey::from_key(key))
                        .into_iter()
                        .flatten(),
                ),
                KeyedState::SexHash(ref mut m) => Box::new(
                    m.swap_remove::<(DfValue, _, _, _, _, _)>(&MakeKey::from_key(key))
                        .into_iter()
                        .flatten(),
                ),
            },
            KeyComparison::Range(range) => {
                macro_rules! remove_range {
                    ($m: expr, $range: expr, $hint: ty) => {
                        Box::new(
                            $m.remove_range(<$hint as MakeKey<DfValue>>::from_range(range))
                                .flat_map(|(_, rows)| rows),
                        )
                    };
                }

                match self.state {
                    KeyedState::SingleBTree(ref mut m) => remove_range!(m, range, DfValue),
                    KeyedState::DoubleBTree(ref mut m) => remove_range!(m, range, (DfValue, _)),
                    KeyedState::TriBTree(ref mut m) => remove_range!(m, range, (DfValue, _, _)),
                    KeyedState::QuadBTree(ref mut m) => {
                        remove_range!(m, range, (DfValue, _, _, _))
                    }
                    KeyedState::QuinBTree(ref mut m) => {
                        remove_range!(m, range, (DfValue, _, _, _, _))
                    }
                    KeyedState::SexBTree(ref mut m) => {
                        remove_range!(m, range, (DfValue, _, _, _, _, _))
                    }
                    KeyedState::MultiBTree(ref mut m, _) => Box::new(
                        m.remove_range::<[DfValue], _>((
                            range.start_bound().map(Vec1::as_slice),
                            range.end_bound().map(Vec1::as_slice),
                        ))
                        .flat_map(|(_, rows)| rows),
                    ),
                    _ => {
                        #[allow(clippy::panic)] // Documented invariant
                        {
                            panic!("mark_hole with a range key called on a HashMap SingleState")
                        }
                    }
                }
            }
        };

        removed
            .filter(|(r, _)| Rc::strong_count(&r.0) == 1)
            .map(|(r, count)| r.deep_size_of() * count)
            .sum()
    }

    pub(super) fn clear(&mut self) {
        self.row_count = 0;
        match self.state {
            KeyedState::AllRows(ref mut rows) => rows.clear(),
            KeyedState::SingleBTree(ref mut map) => map.clear(),
            KeyedState::DoubleBTree(ref mut map) => map.clear(),
            KeyedState::TriBTree(ref mut map) => map.clear(),
            KeyedState::QuadBTree(ref mut map) => map.clear(),
            KeyedState::QuinBTree(ref mut map) => map.clear(),
            KeyedState::SexBTree(ref mut map) => map.clear(),
            KeyedState::MultiBTree(ref mut map, _) => map.clear(),
            KeyedState::SingleHash(ref mut map) => map.clear(),
            KeyedState::DoubleHash(ref mut map) => map.clear(),
            KeyedState::TriHash(ref mut map) => map.clear(),
            KeyedState::QuadHash(ref mut map) => map.clear(),
            KeyedState::QuinHash(ref mut map) => map.clear(),
            KeyedState::SexHash(ref mut map) => map.clear(),
            KeyedState::MultiHash(ref mut map, _) => map.clear(),
        };
    }

    /// Select and evict a random key from `Self::state`. Return the key and evicted rows.
    pub(super) fn evict_random<R: rand::Rng>(
        &mut self,
        rng: &mut R,
    ) -> Option<(Vec<DfValue>, Rows)> {
        self.state
            .evict_with_seed(rng.sample(Uniform::new_inclusive(usize::MIN, usize::MAX).unwrap()))
            .map(|(rows, key)| {
                self.row_count = self.row_count.saturating_sub(rows.len());
                (key, rows)
            })
    }

    /// Evicts a specified key from this state, returning whether any key was present and the
    /// removed rows
    ///
    /// The boolean indicates whether at least one of the keys was present in the state (even if it
    /// had 0 rows). This is useful for distinguishing between:
    /// - A key that doesn't exist in the state → `false`
    /// - A key that exists but has 0 rows → `true`
    /// - A key that exists with rows → `true`
    ///
    /// For range keys, "present" means the range (or part of it) was marked as filled, even if
    /// it contains no actual key entries.
    pub(super) fn evict_keys(&mut self, keys: &[KeyComparison]) -> (bool, Rows) {
        let mut key_was_present = false;
        let evicted = keys
            .iter()
            .flat_map(|k| match k {
                KeyComparison::Equal(equal) => {
                    let result = self.state.evict(equal);
                    // Track if the key was present (even if it had 0 rows)
                    if result.is_some() {
                        key_was_present = true;
                    }
                    Either::Left(
                        result
                            .into_iter()
                            .flat_map(|r| r.into_iter().map(|(r, _)| r)),
                    )
                }
                KeyComparison::Range(range) => {
                    // For ranges, check if any part of the range was marked as filled before evicting.
                    // This allows us to distinguish between:
                    // - A range that was (partially or fully) filled but has no keys
                    // - A range that was never filled at all
                    let range_key = RangeKey::from(range);
                    let lookup_result = self.lookup_range(&range_key);

                    let was_filled = match lookup_result {
                        RangeLookupResult::Some(_) => {
                            // Entire range is fully filled
                            true
                        }
                        RangeLookupResult::Missing(misses) => {
                            // If the misses cover the entire query range, nothing was present.
                            // Otherwise, at least part of the range was present.
                            // Check if there's exactly one miss that matches the entire query range.
                            // Note: misses use Vec<DfValue> while range uses Vec1<DfValue>
                            !(misses.len() == 1
                                && misses[0].0 == range.0.as_ref().map(|v| v.as_vec().clone())
                                && misses[0].1 == range.1.as_ref().map(|v| v.as_vec().clone()))
                        }
                    };

                    let range_rows = self.state.evict_range(range);

                    // Set key_was_present if any part of the range was filled
                    if was_filled {
                        key_was_present = true;
                    }
                    Either::Right(range_rows.into_iter().map(|(r, _)| r))
                }
            })
            .collect::<Rows>();

        // Decrement row_count for evicted rows
        let num_evicted = evicted.len();
        self.row_count = self.row_count.saturating_sub(num_evicted);

        tracing::trace!(
            index = ?self.index,
            num_evicted,
            row_count = self.row_count,
            key_was_present,
            "evicted rows from index"
        );

        (key_was_present, evicted)
    }

    pub(super) fn values<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Rows> + 'a> {
        self.state.values()
    }

    /// Return a reference to this state's Index, which contains the list of columns it's keyed on
    /// and the index type
    pub(super) fn index(&self) -> &Index {
        &self.index
    }

    /// Return this state's index type
    pub(super) fn index_type(&self) -> IndexType {
        self.index().index_type
    }

    /// Return a slice containing the indices of the columns that this index is keyed on
    pub(super) fn columns(&self) -> &[usize] {
        &self.index.columns
    }

    pub(super) fn partial(&self) -> bool {
        self.partial
    }

    pub(super) fn key_count(&self) -> usize {
        self.state.key_count()
    }

    pub(super) fn row_count(&self) -> usize {
        self.row_count
    }

    pub(super) fn is_empty(&self) -> bool {
        self.row_count == 0
    }

    pub(super) fn lookup<'a>(&'a self, key: &PointKey) -> LookupResult<'a> {
        if let Some(rs) = self.state.lookup(key) {
            LookupResult::Some(RecordResult::Borrowed(rs))
        } else if self.partial() {
            // partially materialized, so this is a hole (empty results would be vec![])
            LookupResult::Missing
        } else {
            LookupResult::Some(RecordResult::Owned(vec![]))
        }
    }

    pub(super) fn lookup_range<'a>(&'a self, key: &RangeKey) -> RangeLookupResult<'a> {
        match self.state.lookup_range(key) {
            Ok(rs) => RangeLookupResult::Some(RecordResult::References(rs.collect())),
            Err(misses) if self.partial() => RangeLookupResult::Missing(misses),
            _ => RangeLookupResult::Some(RecordResult::Owned(vec![])),
        }
    }
}

#[cfg(test)]
mod tests {
    use readyset_data::Bound;
    use vec1::vec1;

    use super::*;

    #[test]
    fn mark_filled_point() {
        let mut state = SingleState::new(Index::new(IndexType::BTreeMap, vec![0]), true);
        state.mark_filled(KeyComparison::Equal(vec1![0.into()]));
        assert!(state.lookup(&PointKey::from([0.into()])).is_some());
        assert!(state.is_empty());
    }

    #[test]
    fn mark_filled_range() {
        let mut state = SingleState::new(Index::new(IndexType::BTreeMap, vec![0]), true);
        state.mark_filled(KeyComparison::Range((
            Bound::Included(vec1![0.into()]),
            Bound::Excluded(vec1![5.into()]),
        )));
        assert!(state.lookup(&PointKey::from([0.into()])).is_some());
        assert!(state
            .lookup_range(&RangeKey::from(&(vec1![0.into()]..vec1![5.into()])))
            .is_some());
    }

    mod evict_keys {
        use vec1::vec1;

        use super::*;

        #[test]
        fn equal() {
            let mut state = SingleState::new(Index::new(IndexType::BTreeMap, vec![0]), true);
            let key = KeyComparison::Equal(vec1![0.into()]);
            state.mark_filled(key.clone());
            state.insert_row(vec![0.into(), 1.into()].into());
            let (was_present, _rows) = state.evict_keys(&[key]);
            assert!(was_present, "Key should be present");
            assert!(state.lookup(&PointKey::from([0.into()])).is_missing())
        }

        #[test]
        fn range() {
            let mut state = SingleState::new(Index::new(IndexType::BTreeMap, vec![0]), true);
            let key =
                KeyComparison::from_range(&(vec1![DfValue::from(0)]..vec1![DfValue::from(10)]));
            state.mark_filled(key.clone());
            assert!(state
                .lookup_range(&RangeKey::from(
                    &(vec1![DfValue::from(0)]..vec1![DfValue::from(10)])
                ))
                .is_some());

            state.insert_row(vec![0.into(), 1.into()].into());
            let (was_present, _rows) = state.evict_keys(&[key]);
            assert!(was_present, "Key should be present");
            assert!(state.lookup(&PointKey::from([0.into()])).is_missing());
            assert!(state
                .lookup_range(&RangeKey::from(
                    &(vec1![DfValue::from(0)]..vec1![DfValue::from(10)])
                ))
                .is_missing())
        }

        #[test]
        fn row_count_after_eviction() {
            let mut state = SingleState::new(Index::new(IndexType::HashMap, vec![0]), true);
            state.mark_filled(KeyComparison::Equal(
                Vec1::try_from(vec![1.into()]).unwrap(),
            ));
            state.insert_row(Row::from(vec![1.into(), 2.into()]));
            state.insert_row(Row::from(vec![1.into(), 3.into()]));
            assert_eq!(state.row_count, 2);
            let mut rng = rand::rng();
            state.evict_random(&mut rng);
            assert!(state.is_empty());
        }

        /// Test that evict_keys properly decrements row_count and is_empty works correctly
        /// This is a regression test for the bug where evict_keys didn't decrement row_count,
        /// causing is_empty() to return false even when all keys were evicted.
        /// This bug was particularly problematic for straddled joins with multiple indexes.
        #[test]
        fn evict_keys_decrements_row_count_and_is_empty_works() {
            let mut state = SingleState::new(Index::new(IndexType::HashMap, vec![0]), true);

            // Insert multiple keys with rows
            let key1 = KeyComparison::Equal(vec1![1.into()]);

            state.mark_filled(key1.clone());

            state.insert_row(vec![1.into(), 100.into()].into());
            state.insert_row(vec![1.into(), 101.into()].into());

            // Should have 2 rows total
            assert_eq!(state.row_count, 2);
            assert!(!state.is_empty());

            // Evict key1 (2 rows)
            let (was_present, evicted) = state.evict_keys(&[key1]);
            assert!(was_present, "Key should be present");
            assert_eq!(evicted.len(), 2);
            assert_eq!(state.row_count, 0);
            assert!(state.is_empty());
        }

        /// Test evict_keys with multiple keys in a single call
        #[test]
        fn evict_multiple_keys_at_once() {
            let mut state = SingleState::new(Index::new(IndexType::HashMap, vec![0]), true);

            let key1 = KeyComparison::Equal(vec1![1.into()]);
            let key2 = KeyComparison::Equal(vec1![2.into()]);
            let key3 = KeyComparison::Equal(vec1![3.into()]);

            state.mark_filled(key1.clone());
            state.mark_filled(key2.clone());
            state.mark_filled(key3.clone());

            state.insert_row(vec![1.into(), 100.into()].into());
            state.insert_row(vec![2.into(), 200.into()].into());
            state.insert_row(vec![3.into(), 300.into()].into());

            assert_eq!(state.row_count, 3);

            // Evict all keys at once
            let (was_present, evicted) = state.evict_keys(&[key1, key2, key3]);
            assert!(was_present, "Keys should be present");
            assert_eq!(evicted.len(), 3);
            assert_eq!(state.row_count, 0);
            assert!(
                state.is_empty(),
                "is_empty() should return true after evicting all keys at once"
            );
        }

        /// Test with BTreeMap (range eviction support)
        #[test]
        fn evict_keys_btree_decrements_row_count() {
            let mut state = SingleState::new(Index::new(IndexType::BTreeMap, vec![0]), true);

            // Fill a range
            let range = KeyComparison::from_range(&(vec1![1.into()]..vec1![10.into()]));
            state.mark_filled(range.clone());

            // Insert rows in the range
            state.insert_row(vec![1.into(), 100.into()].into());
            state.insert_row(vec![5.into(), 500.into()].into());
            state.insert_row(vec![9.into(), 900.into()].into());

            assert_eq!(state.row_count, 3);
            assert!(!state.is_empty());

            // Evict the entire range
            let (was_present, evicted) = state.evict_keys(&[range]);
            assert!(was_present, "Range should be present");
            assert_eq!(evicted.len(), 3);
            assert_eq!(state.row_count, 0);
            assert!(
                state.is_empty(),
                "is_empty() should return true after range eviction"
            );
        }

        /// Test partial eviction doesn't incorrectly report empty
        #[test]
        fn evict_keys_partial_not_empty() {
            let mut state = SingleState::new(Index::new(IndexType::HashMap, vec![0]), true);

            let key1 = KeyComparison::Equal(vec1![1.into()]);
            let key2 = KeyComparison::Equal(vec1![2.into()]);

            state.mark_filled(key1.clone());
            state.mark_filled(key2.clone());

            state.insert_row(vec![1.into(), 100.into()].into());
            state.insert_row(vec![2.into(), 200.into()].into());

            assert_eq!(state.row_count, 2);
            assert!(!state.is_empty());

            // Evict only one key
            let (was_present, _rows) = state.evict_keys(&[key1]);
            assert!(was_present, "Key should be present");
            assert_eq!(state.row_count, 1);
            assert!(
                !state.is_empty(),
                "State should not be empty with one key remaining"
            );

            // Evict the last key
            let (was_present, _rows) = state.evict_keys(&[key2]);
            assert!(was_present, "Key should be present");
            assert_eq!(state.row_count, 0);
            assert!(
                state.is_empty(),
                "State should be empty after all keys evicted"
            );
        }

        /// Test that evict_keys returns true for keys that exist but have 0 rows
        /// This is important for tracking whether a key was filled but empty vs never existed
        #[test]
        fn evict_keys_distinguishes_empty_key_from_missing_key() {
            let mut state = SingleState::new(Index::new(IndexType::HashMap, vec![0]), true);

            // Case 1: Key is marked as filled but has no rows
            let key_filled_empty = KeyComparison::Equal(vec1![1.into()]);
            state.mark_filled(key_filled_empty.clone());
            // Don't insert any rows for this key

            let (was_present, rows) = state.evict_keys(&[key_filled_empty]);
            assert!(
                was_present,
                "Key that was filled but has 0 rows should return true"
            );
            assert_eq!(rows.len(), 0, "Should return 0 rows");

            // Case 2: Key was never filled (doesn't exist in state)
            let key_never_filled = KeyComparison::Equal(vec1![2.into()]);
            let (was_present, rows) = state.evict_keys(&[key_never_filled]);
            assert!(
                !was_present,
                "Key that was never filled should return false"
            );
            assert_eq!(rows.len(), 0, "Should return 0 rows");

            // Case 3: Key with rows
            let key_with_rows = KeyComparison::Equal(vec1![3.into()]);
            state.mark_filled(key_with_rows.clone());
            state.insert_row(vec![3.into(), 300.into()].into());

            let (was_present, rows) = state.evict_keys(&[key_with_rows]);
            assert!(was_present, "Key with rows should return true");
            assert_eq!(rows.len(), 1, "Should return 1 row");
        }

        /// Test evict_keys with BTreeMap for point keys (Equal)
        /// This ensures the boolean return value works correctly for BTreeMap indices
        #[test]
        fn evict_keys_btree_point_keys() {
            let mut state = SingleState::new(Index::new(IndexType::BTreeMap, vec![0]), true);

            // Case 1: Key is marked as filled but has no rows
            let key_filled_empty = KeyComparison::Equal(vec1![1.into()]);
            state.mark_filled(key_filled_empty.clone());
            // Don't insert any rows for this key

            let (was_present, rows) = state.evict_keys(&[key_filled_empty]);
            assert!(
                was_present,
                "BTreeMap: Key that was filled but has 0 rows should return true"
            );
            assert_eq!(rows.len(), 0, "Should return 0 rows");

            // Case 2: Key was never filled (doesn't exist in state)
            let key_never_filled = KeyComparison::Equal(vec1![2.into()]);
            let (was_present, rows) = state.evict_keys(&[key_never_filled]);
            assert!(
                !was_present,
                "BTreeMap: Key that was never filled should return false"
            );
            assert_eq!(rows.len(), 0, "Should return 0 rows");

            // Case 3: Key with rows
            let key_with_rows = KeyComparison::Equal(vec1![3.into()]);
            state.mark_filled(key_with_rows.clone());
            state.insert_row(vec![3.into(), 300.into()].into());
            state.insert_row(vec![3.into(), 301.into()].into());

            let (was_present, rows) = state.evict_keys(&[key_with_rows]);
            assert!(was_present, "BTreeMap: Key with rows should return true");
            assert_eq!(rows.len(), 2, "Should return 2 rows");
            assert_eq!(state.row_count, 0, "Row count should be 0 after eviction");
        }

        /// Test evict_keys with BTreeMap for range keys
        #[test]
        fn evict_keys_btree_range_keys() {
            let mut state = SingleState::new(Index::new(IndexType::BTreeMap, vec![0]), true);

            // Case 1: Range is marked as filled but has no rows
            let range_filled_empty =
                KeyComparison::from_range(&(vec1![10.into()]..vec1![20.into()]));
            let included_small_range_1 =
                KeyComparison::from_range(&(vec1![10.into()]..vec1![15.into()]));
            let included_small_range_2 =
                KeyComparison::from_range(&(vec1![16.into()]..vec1![20.into()]));
            let excluded_small_range_1 =
                KeyComparison::from_range(&(vec1![14.into()]..vec1![20.into()]));
            state.mark_filled(range_filled_empty.clone());
            // Don't insert any rows in this range

            let (was_present, rows) = state.evict_keys(&[included_small_range_1]);
            assert!(
                was_present,
                "BTreeMap: Range that was filled but has 0 keys should return true"
            );
            assert_eq!(rows.len(), 0, "Should return 0 rows");

            let (was_present, rows) = state.evict_keys(&[excluded_small_range_1]);
            assert!(
                was_present,
                "BTreeMap: the range that overlaps with filled and holes should return true"
            );
            assert_eq!(rows.len(), 0, "Should return 0 rows");

            let (was_present, rows) = state.evict_keys(&[included_small_range_2]);
            assert!(
                !was_present,
                "BTreeMap: Range was completely evicted from previous calls. Should return false"
            );
            assert_eq!(rows.len(), 0, "Should return 0 rows");

            // Case 2: Range was never filled (doesn't exist in state)
            let range_never_filled =
                KeyComparison::from_range(&(vec1![30.into()]..vec1![40.into()]));
            let (was_present, rows) = state.evict_keys(&[range_never_filled]);
            assert!(
                !was_present,
                "BTreeMap: Range that was never filled should return false"
            );
            assert_eq!(rows.len(), 0, "Should return 0 rows");

            // Case 3: Range with rows
            let range_with_rows = KeyComparison::from_range(&(vec1![50.into()]..vec1![60.into()]));
            state.mark_filled(range_with_rows.clone());
            state.insert_row(vec![50.into(), 500.into()].into());
            state.insert_row(vec![55.into(), 550.into()].into());
            state.insert_row(vec![59.into(), 590.into()].into());

            let (was_present, rows) = state.evict_keys(&[range_with_rows]);
            assert!(was_present, "BTreeMap: Range with rows should return true");
            assert_eq!(rows.len(), 3, "Should return 3 rows");
            assert_eq!(state.row_count, 0, "Row count should be 0 after eviction");
        }

        /// Test evict_keys with BTreeMap for multiple keys at once
        #[test]
        fn evict_keys_btree_multiple_keys() {
            let mut state = SingleState::new(Index::new(IndexType::BTreeMap, vec![0]), true);

            // Setup: Create mix of filled/unfilled keys and ranges
            let key1 = KeyComparison::Equal(vec1![1.into()]);
            let key2 = KeyComparison::Equal(vec1![2.into()]);
            let key3 = KeyComparison::Equal(vec1![3.into()]);
            let range1 = KeyComparison::from_range(&(vec1![10.into()]..vec1![15.into()]));

            state.mark_filled(key1.clone());
            state.insert_row(vec![1.into(), 100.into()].into());

            state.mark_filled(key2.clone());
            // key2 has no rows

            // key3 is never filled

            state.mark_filled(range1.clone());
            state.insert_row(vec![10.into(), 1000.into()].into());
            state.insert_row(vec![12.into(), 1200.into()].into());

            // Evict all at once
            let (was_present, rows) = state.evict_keys(&[key1, key2, key3.clone(), range1]);
            assert!(
                was_present,
                "BTreeMap: Should return true if at least one key was present"
            );
            assert_eq!(
                rows.len(),
                3,
                "Should return 3 rows total (1 from key1, 2 from range1)"
            );
            assert_eq!(state.row_count, 0);

            // Verify all keys are now missing
            assert!(state.lookup(&PointKey::from([1.into()])).is_missing());
            assert!(state.lookup(&PointKey::from([2.into()])).is_missing());
        }

        /// Test evict_keys with BTreeMap multi-column index
        #[test]
        fn evict_keys_btree_multi_column() {
            let mut state = SingleState::new(Index::new(IndexType::BTreeMap, vec![0, 1]), true);

            // Case 1: Multi-column key that's filled but empty
            let key_empty = KeyComparison::Equal(vec1![1.into(), 10.into()]);
            state.mark_filled(key_empty.clone());

            let (was_present, rows) = state.evict_keys(&[key_empty]);
            assert!(
                was_present,
                "BTreeMap multi-column: Empty key should return true"
            );
            assert_eq!(rows.len(), 0);

            // Case 2: Multi-column key with rows
            let key_with_rows = KeyComparison::Equal(vec1![2.into(), 20.into()]);
            state.mark_filled(key_with_rows.clone());
            state.insert_row(vec![2.into(), 20.into(), 200.into()].into());
            state.insert_row(vec![2.into(), 20.into(), 201.into()].into());

            let (was_present, rows) = state.evict_keys(&[key_with_rows]);
            assert!(
                was_present,
                "BTreeMap multi-column: Key with rows should return true"
            );
            assert_eq!(rows.len(), 2);

            // Case 3: Multi-column key never filled
            let key_missing = KeyComparison::Equal(vec1![3.into(), 30.into()]);
            let (was_present, rows) = state.evict_keys(&[key_missing]);
            assert!(
                !was_present,
                "BTreeMap multi-column: Missing key should return false"
            );
            assert_eq!(rows.len(), 0);
        }

        /// Test evict_keys with BTreeMap - edge case of evicting only unfilled keys
        #[test]
        fn evict_keys_btree_only_missing_keys() {
            let mut state = SingleState::new(Index::new(IndexType::BTreeMap, vec![0]), true);

            // Try to evict multiple keys that were never filled
            let key1 = KeyComparison::Equal(vec1![1.into()]);
            let key2 = KeyComparison::Equal(vec1![2.into()]);
            let key3 = KeyComparison::Equal(vec1![3.into()]);

            let (was_present, rows) = state.evict_keys(&[key1, key2, key3]);
            assert!(
                !was_present,
                "BTreeMap: Should return false when no keys were present"
            );
            assert_eq!(rows.len(), 0, "Should return 0 rows");
        }
    }
}
